package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

const (
	IPPROTO_OSPF   = 89
	OSPFVersion    = 2
	OSPFHelloType  = 1
	OSPFHeaderLen  = 24
	HelloPacketLen = 20
	IPHeaderLen    = 20
	BufferSize     = 65535
)

type OSPFServer struct {
	fd         int
	ifaceName  string
	outputFile *os.File
	buffer     []byte
	stopChan   chan struct{}
	wg         sync.WaitGroup
	mu         sync.Mutex
}

type PacketStats struct {
	TotalPackets uint64
	OSPFPackets  uint64
	HelloPackets uint64
	LastReceived time.Time
	mu           sync.RWMutex
}

// 创建新的OSPF服务器
func NewOSPFServer(ifaceName, outputPath string) (*OSPFServer, error) {
	// 创建原始socket
	fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_RAW, IPPROTO_OSPF)
	if err != nil {
		return nil, fmt.Errorf("创建socket失败: %v", err)
	}

	// 设置socket选项
	err = syscall.SetsockoptInt(fd, syscall.IPPROTO_IP, syscall.IP_HDRINCL, 1)
	if err != nil {
		syscall.Close(fd)
		return nil, fmt.Errorf("设置socket选项失败: %v", err)
	}

	// 打开输出文件
	file, err := os.OpenFile(outputPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		syscall.Close(fd)
		return nil, fmt.Errorf("打开输出文件失败: %v", err)
	}

	return &OSPFServer{
		fd:         fd,
		ifaceName:  ifaceName,
		outputFile: file,
		buffer:     make([]byte, BufferSize),
		stopChan:   make(chan struct{}),
	}, nil
}

// 启动服务器
func (s *OSPFServer) Start() error {
	log.Printf("OSPF服务器启动，监听接口: %s", s.ifaceName)

	s.wg.Add(1)
	go s.receivePackets()

	return nil
}

// 停止服务器
func (s *OSPFServer) Stop() {
	log.Println("正在停止OSPF服务器...")
	close(s.stopChan)
	s.wg.Wait()

	if s.outputFile != nil {
		s.outputFile.Close()
	}
	if s.fd != 0 {
		syscall.Close(s.fd)
	}
	log.Println("OSPF服务器已停止")
}

// 接收数据包
func (s *OSPFServer) receivePackets() {
	defer s.wg.Done()

	var stats PacketStats

	for {
		select {
		case <-s.stopChan:
			return
		default:
			n, addr, err := syscall.Recvfrom(s.fd, s.buffer, 0)
			if err != nil {
				log.Printf("接收数据失败: %v", err)
				continue
			}

			stats.mu.Lock()
			stats.TotalPackets++
			stats.LastReceived = time.Now()
			stats.mu.Unlock()

			if n < IPHeaderLen+OSPFHeaderLen {
				continue
			}

			// 解析IP头
			ipHeader := s.buffer[:IPHeaderLen]
			if ipHeader[9] != IPPROTO_OSPF {
				continue
			}

			stats.mu.Lock()
			stats.OSPFPackets++
			stats.mu.Unlock()

			// 解析OSPF头
			ospfHeader := s.buffer[IPHeaderLen : IPHeaderLen+OSPFHeaderLen]
			if ospfHeader[0] != OSPFVersion || ospfHeader[1] != OSPFHelloType {
				continue
			}

			stats.mu.Lock()
			stats.HelloPackets++
			stats.mu.Unlock()

			// 提取源IP地址
			srcIP := net.IPv4(addr.(*syscall.SockaddrInet4).Addr[0],
				addr.(*syscall.SockaddrInet4).Addr[1],
				addr.(*syscall.SockaddrInet4).Addr[2],
				addr.(*syscall.SockaddrInet4).Addr[3]).String()

			// 提取Dead Interval字段中的数据
			if n >= IPHeaderLen+OSPFHeaderLen+HelloPacketLen {
				deadInterval := s.buffer[IPHeaderLen+OSPFHeaderLen+8 : IPHeaderLen+OSPFHeaderLen+12]
				s.writeToFile(srcIP, deadInterval)
			}

			// 每10秒打印一次统计信息
			if time.Since(stats.LastReceived) >= 10*time.Second {
				stats.mu.RLock()
				log.Printf("统计信息 - 总数据包: %d, OSPF包: %d, Hello包: %d",
					stats.TotalPackets, stats.OSPFPackets, stats.HelloPackets)
				stats.mu.RUnlock()
			}
		}
	}
}

// 写入数据到文件
func (s *OSPFServer) writeToFile(srcIP string, data []byte) {
	s.mu.Lock()
	defer s.mu.Unlock()

	timestamp := time.Now().Format("2006-01-02 15:04:05.000")
	entry := fmt.Sprintf("[%s] 从 %s 接收到数据: %v\n", timestamp, srcIP, data)

	if _, err := s.outputFile.WriteString(entry); err != nil {
		log.Printf("写入文件失败: %v", err)
		return
	}

	s.outputFile.Sync()
}

func main() {
	// 解析命令行参数
	ifaceName := flag.String("i", "", "网络接口名称")
	outputPath := flag.String("o", "ospf_data.txt", "输出文件路径")
	flag.Parse()

	if *ifaceName == "" {
		log.Fatal("请指定网络接口名称")
	}

	// 创建服务器
	server, err := NewOSPFServer(*ifaceName, *outputPath)
	if err != nil {
		log.Fatalf("创建服务器失败: %v", err)
	}

	// 启动服务器
	if err := server.Start(); err != nil {
		log.Fatalf("启动服务器失败: %v", err)
	}

	// 处理中断信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	// 等待中断信号
	sig := <-sigChan
	log.Printf("接收到信号 %v, 正在关闭...", sig)

	// 优雅关闭
	server.Stop()
}
